# -*- codeing = utf-8 -*-
# @Time : 2022/11/7 16:21
# @Author : xiaow
# @File : videosocket.py
# @Software : PyCharm
# 获取图片数据，每帧传递给前台从而做成视频


import asyncio
import time

import websockets
import base64
import cv2
import numpy as np
from concurrent.futures.thread import ThreadPoolExecutor
from multiprocessing import Lock
import mysqldao
from websockets.exceptions import ConnectionClosedError

lock = Lock()
outlock = Lock()
infopics = []
videopics = []
startflag = False
videoname = 'xw2.avi'
# 线程池
pool = ThreadPoolExecutor(max_workers=50)
task_list = []
startflag = False


# 存储图片和信息的类
class MImg:
    #  t 时间  data 图像数据
    def __init__(self, t, data, code=-1, data2="", info=""):
        self.t = t
        self.data = data
        # 状态码 -1 没有检测图片  0正常  1裂纹
        self.code = code
        # 检测后的图片
        self.data2 = data2
        # 监测信息
        self.info = info


# 输出视频
def outVideo(out, frame):
    global startflag
    if startflag:
        out.write(frame)


# 处理图像 这里的img就是三通道图像数据
def dealwithpic(img, lock):
    global infopics
    # 操作图像start
    # 操作end
    lock.acquire()
    infopics.append(MImg(time.time(), img, 0, img, '检测正常'))
    infopics = sorted(infopics, key=lambda e: e.t)
    lock.release()


# 清理线程池中执行完的线程
def clearpool():
    while (True):
        if len(task_list) == 0:
            break
        for i in task_list:
            if i.done():
                task_list.remove(i)


# 向服务器端实时发送视频截图
async def send_msg(websocket):
    global startflag
    if startflag:
        capture = cv2.VideoCapture('rtsp://127.0.0.1/test')
        if not capture.isOpened():
            print('quit')
            quit()
        ret, frame = capture.read()
        # print(1)
        fps = 30
        fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', '2')
        width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
        out = cv2.VideoWriter('D:\\output.avi', fourcc, fps, (width, height))
        all = 0
        global lock, videopics
        while True:
            # 存储帧数
            all += 1

            # 开启新线程存储视频
            task_list.append(pool.submit(outVideo, out, frame))
            frame = cv2.flip(frame, -1)
            frame = cv2.flip(frame, 0)
            result, imgencode = cv2.imencode('.jpg', frame)
            data = np.array(imgencode)
            data = cv2.imdecode(data, cv2.IMREAD_COLOR)
            img_encode = cv2.imencode('.jpg', data)[1]
            base64_data = base64.b64encode(img_encode)
            base64_data_str = str(base64_data, encoding="utf-8")
            # code!=-1就代表有检测完成的图片
            code = -1
            info = ''
            if all % 30 == 0:
                task_list.append(pool.submit(dealwithpic, data, lock))
                all = 0
            lock.acquire()
            if infopics.__len__() > 0:
                a = infopics.pop()
                img_encode2 = cv2.imencode('.jpg', a.data2)[1]
                base64_data2 = base64.b64encode(img_encode2)
                base64_data_str2 = str(base64_data2, encoding="utf-8")
                code = a.code
                info = a.info
            else:
                base64_data_str2 = ""
            result = '{"data":"%s","code":"%s","data2":"%s","info":"%s"}' % ("data:image/jpg;base64," + base64_data_str,
                                                                             code.__str__(),
                                                                             "data:image/jpg;base64," + base64_data_str2,
                                                                             info
                                                                             )
            lock.release()
            try:
                await websocket.send(result)
            except ConnectionClosedError as e:
                break
            ret, frame = capture.read()
        # 添加至数据库
        mysqldao.insert(db='robot', name='测试管道', videopath='output.avi', imgspath='')
        out.release()


async def serverRecv(websocket):
    global startflag, videoname, infopics, videopics
    # 设置存储名
    videoname = "D:\\new.avi"
    while True:
        try:
            startflag = False
            # 一直循环，等待连接
            recv_text = await websocket.recv()
            # 接收到连接之后执行
            startflag = True
        except ConnectionClosedError as e:
            pass
        # 断开连接后，send_msg仍然会继续执行
        await send_msg(websocket)


if __name__ == '__main__':
    task_list.append(pool.submit(clearpool))
    server = websockets.serve(serverRecv, "127.0.0.1", "8888")
    asyncio.get_event_loop().run_until_complete(server)
    asyncio.get_event_loop().run_forever()
